# -*- coding: utf-8 -*-
import pandas as pd
import sys
import os

# --- 1. 从命令行获取参数 ---
if len(sys.argv) < 3:
    print("错误：参数不足。")
    print("用法: python3 02_preprocess_data.py <股票代码> <frequency>")
    sys.exit(1)

stock_code = sys.argv[1]
frequency = sys.argv[2]

# --- 2. 定义输入输出文件名 ---
raw_filename = f'output/data/stock_data_{stock_code}_{frequency}_raw.csv'
processed_filename = f'output/data/stock_data_{stock_code}_{frequency}.csv'

# --- 3. 加载原始数据 ---
try:
    print(f"开始加载原始数据文件: {raw_filename}...")
    df = pd.read_csv(raw_filename)
    print("加载成功。")
except FileNotFoundError:
    print(f"错误: 找不到原始数据文件 {raw_filename}。")
    sys.exit(1)

# --- 4. 数据预处理 ---
print("开始进行数据预处理...")
# '日期' 列在分钟线数据里可能叫 '时间'，我们统一处理
date_col = '日期'
if '时间' in df.columns:
    date_col = '时间'

columns_to_keep = [date_col, '股票代码', '开盘', '收盘', '最高', '最低', '成交量', '成交额', '振幅', '涨跌幅', '涨跌额', '换手率']
df_processed = df[columns_to_keep]
df_processed.rename(columns={date_col: '日期'}, inplace=True) # 统一列名为'日期'
df_processed.dropna(inplace=True)
print("预处理完成。")

# --- 5. 保存处理后的数据 ---
# 确保output/data目录存在
os.makedirs('output/data', exist_ok=True)

df_processed.to_csv(processed_filename, index=False)
print(f"处理后的数据已保存至: {processed_filename}")